import os
import numpy as np
from skimage.io import imread
from skimage.transform import resize
from skimage.feature import hog
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
import joblib


def load_dataset(dataset_path):
    """加载数据集并生成路径-标签对"""
    image_paths = []
    labels = []
    # 按文件夹名称排序保证一致性
    for label, class_folder in enumerate(sorted(os.listdir(dataset_path))):
        class_dir = os.path.join(dataset_path, class_folder)
        if os.path.isdir(class_dir):
            # 收集当前类别的所有图像路径
            for img_file in os.listdir(class_dir):
                img_path = os.path.join(class_dir, img_file)
                image_paths.append(img_path)
                labels.append(label)  # 直接使用排序序号作为标签
    return image_paths, labels


def extract_features(image_paths):
    """批量提取HOG特征"""
    features = []
    for path in image_paths:
        img = imread(path, as_gray=True)
        img = resize(img, (64, 64))
        features.append(hog(img,
                            orientations=9,
                            pixels_per_cell=(8, 8),
                            cells_per_block=(2, 2),
                            visualize=False))
    return np.array(features)


# 1. 加载数据集
# 数据集路径
train_dataset_path = r"F:\人工智能教材编写\traffic_sign\train"
train_paths, train_labels = load_dataset(train_dataset_path)

# 2. 提取特征
X = extract_features(train_paths)
y = np.array(train_labels)

# 3. 划分训练验证集
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

# 4. 训练并保存模型
knn = KNeighborsClassifier(n_neighbors=5)
knn.fit(X_train, y_train)
joblib.dump(knn, "knn_model.pkl")

# 5. 验证模型
print(f"验证准确率: {knn.score(X_val, y_val):.2%}")
